package ocr

import (
	"log"
	"path/filepath"
	"time"

	"gocv.io/x/gocv"
)

type TextClassifier struct {
	*PaddleModel
	batchNum int
	thresh   float64
	shape    []int
	labels   []string
}

type ClsResult struct {
	Score float32
	Label int64
}

func NewTextClassifier(modelSource string, config *Config) *TextClassifier {
	shapes := config.ClsImageShape // []int{3, 48, 192}
	// if v, ok := args["cls_image_shape"]; ok {
	// 	for i, s := range v.([]interface{}) {
	// 		shapes[i] = s.(int)
	// 	}
	// }
	cls := &TextClassifier{
		PaddleModel: NewPaddleModel("cls", config),
		batchNum:    config.ClsBatchNum, //getInt(args, "cls_batch_num", 30),
		thresh:      config.ClsThresh,   //getFloat64(args, "cls_thresh", 0.9),
		shape:       shapes,
	}
	if checkModelExists(modelSource) {
		modelSource, _ = downloadModel(filepath.Join(config.ModleDir), modelSource)
	} else {
		log.Panicf("cls model path: %v not exist! Please check!", modelSource)
	}
	cls.LoadModel(modelSource)
	return cls
}

func (cls *TextClassifier) Run(imgs []gocv.Mat) []gocv.Mat {
	batch := cls.batchNum
	var clsTime int64 = 0
	clsout := make([]ClsResult, len(imgs))
	srcimgs := make([]gocv.Mat, len(imgs))
	c, h, w := cls.shape[0], cls.shape[1], cls.shape[2]
	for i := 0; i < len(imgs); i += batch {
		j := i + batch
		if len(imgs) < j {
			j = len(imgs)
		}

		normImgs := make([]float32, (j-i)*c*h*w)
		for k := i; k < j; k++ {
			tmp := gocv.NewMat()
			imgs[k].CopyTo(&tmp)
			srcimgs[k] = tmp
			img := clsResize(imgs[k], cls.shape)
			data := normPermute(img, []float32{0.5, 0.5, 0.5}, []float32{0.5, 0.5, 0.5}, 255.0)
			copy(normImgs[(k-i)*c*h*w:], data)
		}

		st := time.Now()
		cls.inputs[cls.inNames[0]].CopyFromCpu(normImgs)
		cls.inputs[cls.inNames[0]].Reshape([]int32{int32(j - i), int32(c), int32(h), int32(w)})

		// cls.predictor.SetZeroCopyInput(cls.input)
		// cls.predictor.ZeroCopyRun()
		// cls.predictor.GetZeroCopyOutput(cls.outputs[0])
		// cls.predictor.GetZeroCopyOutput(cls.outputs[1])

		cls.predictor.Run()

		outData := make([]float32, numElements(cls.outputs[cls.outNames[0]].Shape()))
		cls.outputs[cls.outNames[0]].CopyToCpu(outData)

		var probout [][]float32
		var labelout []int64
		if len(cls.outputs[cls.outNames[0]].Shape()) == 2 {
			cls.inputs[cls.outNames[0]].CopyToCpu(probout)
		} else {
			cls.inputs[cls.outNames[0]].CopyToCpu(labelout)
		}

		if len(cls.outputs[cls.outNames[1]].Shape()) == 2 {
			cls.inputs[cls.outNames[1]].CopyToCpu(probout)
		} else {
			cls.inputs[cls.outNames[1]].CopyToCpu(labelout)
		}
		clsTime += int64(time.Since(st).Milliseconds())

		for no, label := range labelout {
			score := probout[no][label]
			clsout[i+no] = ClsResult{
				Score: score,
				Label: label,
			}

			if label%2 == 1 && float64(score) > cls.thresh {
				gocv.Rotate(srcimgs[i+no], &srcimgs[i+no], gocv.Rotate180Clockwise)
			}
		}
	}
	log.Println("cls num: ", len(clsout), ", cls time elapse: ", clsTime, "ms")
	return srcimgs
}
